import os
import time

from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys


def change_sk(driver, code):
    """
    切换到指定公司

    参数:
        driver: selenium webdriver实例
        code: 公司代码

    功能:
        1. 切换到默认内容
        2. 等待页面加载完成
        3. 获取当前公司名称
        4. 尝试切换公司
        5. 验证是否切换成功
    """
    driver.switch_to.default_content()  # 切换到默认内容

    # 等待页面加载完成
    while True:
        try:
            driver.find_element(By.CLASS_NAME, 'area-component')  # 查找区域组件
            break
        except:
            time.sleep(2)  # 等待2秒

    # 获取当前公司名称
    old_name = driver.find_element(By.XPATH,
                                   '/html/body/div[1]/div/div[2]/section/header/div/div[2]/div[3]/span/span/div/div/span').text
    print('上一次所在公司为:', old_name)

    # 尝试最多3次切换公司
    for i in range(3):
        driver.switch_to.default_content()  # 每次尝试前切换到默认内容

        # 获取当前公司名称
        while True:
            try:
                now_name = driver.find_element(By.XPATH,
                                               '/html/body/div[1]/div/div[2]/section/header/div/div[2]/div[3]/span/span/div/div/span').text
                break
            except:
                time.sleep(2)  # 等待2秒

        # 检查是否已切换成功
        if now_name != old_name:
            print('当前所在公司为:', now_name)
            break
        else:
            # 点击公司选择器
            while True:
                try:
                    driver.find_element(By.XPATH,
                                        '//*[@id="app"]/div/div[2]/section/header/div/div[2]/div[3]/span/span/div/div/span').click()
                    time.sleep(2)
                    break
                except:
                    time.sleep(2)

            time.sleep(2)

            # 尝试输入公司代码并切换
            try:
                # 清空输入框
                driver.find_element(By.CLASS_NAME, 'el-input__inner').clear()
                time.sleep(1)

                # 输入公司代码并回车
                driver.find_element(By.CLASS_NAME, 'el-input__inner').send_keys(code, Keys.ENTER)
                time.sleep(4)

                # 查找并点击对应公司
                while True:
                    try:
                        driver.find_element_by_xpath(f"//*[text()='{code}']").click()
                        break
                    except:
                        time.sleep(2)

                time.sleep(5)
            except Exception as e:
                print(e)  # 打印异常信息


def wait_download(before_download, download_dir):
    # before_download = set(os.listdir(download_dir))
    while True:
        # 检查是否有.crdownload临时文件
        temp_files = [f for f in os.listdir(download_dir) if f.endswith('.crdownload')]
        if not temp_files:
            # 检查文件大小是否稳定
            after_download = set(os.listdir(download_dir))
            new_files = after_download - before_download
            if new_files:
                file_path = os.path.join(download_dir, next(iter(new_files)))
                size1 = os.path.getsize(file_path)
                time.sleep(1)
                size2 = os.path.getsize(file_path)
                if size1 == size2:
                    break